"""加密工具模块

提供 token 加密/解密功能，由调用方决定是否使用
"""

import base64
import hashlib
from typing import Optional, Union

try:
    from cryptography.fernet import Fernet
    from cryptography.hazmat.primitives import hashes
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
    _CRYPTOGRAPHY_AVAILABLE = True
except ImportError:
    _CRYPTOGRAPHY_AVAILABLE = False


def _derive_key(password: str, salt: bytes) -> bytes:
    """从密码派生加密密钥"""
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
    )
    return base64.urlsafe_b64encode(kdf.derive(password.encode()))


def encrypt_token(token: str, password: str) -> str:
    """加密 token
    
    Args:
        token: 要加密的 token
        password: 加密密码
        
    Returns:
        str: 加密后的 token (Base64 编码)
        
    Raises:
        ImportError: 如果 cryptography 库未安装
    """
    if not _CRYPTOGRAPHY_AVAILABLE:
        raise ImportError("cryptography 库未安装，无法使用加密功能。请安装: pip install cryptography")
    
    # 生成随机盐
    salt = base64.urlsafe_b64encode(hashlib.sha256(password.encode()).digest())[:16]
    
    # 派生密钥
    key = _derive_key(password, salt)
    
    # 加密
    f = Fernet(key)
    encrypted_token = f.encrypt(token.encode())
    
    # 组合盐和加密数据
    result = salt + b":" + encrypted_token
    
    # 返回 Base64 编码的结果
    return base64.urlsafe_b64encode(result).decode()


def decrypt_token(encrypted_token: str, password: str) -> str:
    """解密 token
    
    Args:
        encrypted_token: 加密的 token (Base64 编码)
        password: 解密密码
        
    Returns:
        str: 解密后的 token
        
    Raises:
        ImportError: 如果 cryptography 库未安装
        ValueError: 如果解密失败
    """
    if not _CRYPTOGRAPHY_AVAILABLE:
        raise ImportError("cryptography 库未安装，无法使用加密功能。请安装: pip install cryptography")
    
    try:
        # Base64 解码
        data = base64.urlsafe_b64decode(encrypted_token.encode())
        
        # 分离盐和加密数据
        salt, encrypted_data = data.split(b":", 1)
        
        # 派生密钥
        key = _derive_key(password, salt)
        
        # 解密
        f = Fernet(key)
        token = f.decrypt(encrypted_data)
        
        return token.decode()
    except Exception as e:
        raise ValueError(f"解密失败: {str(e)}")


def is_crypto_available() -> bool:
    """检查加密功能是否可用"""
    return _CRYPTOGRAPHY_AVAILABLE